module gw_front

  !
  ! This module handles gravity waves from frontal sources, and was extracted
  ! from gw_drag in May 2013.
  !

  use gw_utils,  only: r8
  use gw_common, only: GWBand, pi, pver


  implicit none

  private

  save

  public CMSourceDesc
  public flat_cm_desc
  public gaussian_cm_desc
  public gw_cm_src

  ! Tuneable settings.
  type CMSourceDesc
    ! Source level.
    integer ksrc
    ! Level at which to check whether the frontogenesis function is above
    ! the critical threshold.
    integer kfront
    ! Frontogenesis function critical threshold.
    real(r8) :: frontgfc = huge(1._r8)
    ! The stress spectrum to produce at the source level.
    real(r8), allocatable :: src_tau(:)
  end type CMSourceDesc

contains

  ! Create a flat profile to be launched (all wavenumbers have the same
  ! source strength, except that l=0 is excluded).
  function flat_cm_desc(band, ksrc, kfront, frontgfc, taubgnd) result(desc)
    ! Wavelengths triggered by frontogenesis.
    type(GWBand), intent(in) :: band
    ! The following are used to set the corresponding object components.
    integer, intent(in) :: ksrc
    integer, intent(in) :: kfront
    real(r8), intent(in) :: frontgfc
    ! Amount of stress to launch from each wavelength.
    real(r8), intent(in) :: taubgnd

    type(CMSourceDesc) desc

    desc%ksrc = ksrc
    desc%kfront = kfront
    desc%frontgfc = frontgfc

    allocate(desc%src_tau(-band%ngwv:band%ngwv))
    desc%src_tau = taubgnd

    ! Prohibit wavenumber 0.
    desc%src_tau(0) = 0

  end function flat_cm_desc

  ! Create a source tau profile that is a gaussian over wavenumbers (l=0 is excluded).
  function gaussian_cm_desc(band, ksrc, kfront, frontgfc, height, width) result(desc)

    use shr_spfn_mod, only: erfc => shr_spfn_erfc

    ! Wavelengths triggered by frontogenesis.
    type(GWBand), intent(in) :: band
    ! The following are used to set the corresponding object components.
    integer, intent(in) :: ksrc
    integer, intent(in) :: kfront
    real(r8), intent(in) :: frontgfc
    ! Parameters of gaussian.
    real(r8), intent(in) :: height
    real(r8), intent(in) :: width

    type(CMSourceDesc) desc

    ! Bounds used to average bins of the gaussian.
    real(r8) gaussian_bounds(2*band%ngwv+2)

    ! Wavenumber index.
    integer l

    desc%ksrc = ksrc
    desc%kfront = kfront
    desc%frontgfc = frontgfc

    allocate(desc%src_tau(-band%ngwv:band%ngwv))

    ! Find the boundaries of each bin.
    gaussian_bounds(:2*band%ngwv+1) = band%cref - 0.5_r8*band%dc
    gaussian_bounds(2*band%ngwv+2) = band%cref(band%ngwv) + 0.5_r8*band%dc

    ! Integral of the gaussian at bin interfaces (from the point to
    ! positive infinity).
    gaussian_bounds = &
        [( erfc(gaussian_bounds(l)/width)*height*width*sqrt(pi)/2._r8, &
        l = 1, 2*band%ngwv+2 )]

    ! Get average in each bin using integral from right to left side.
    desc%src_tau = &
        (gaussian_bounds(:2*band%ngwv+1) - gaussian_bounds(2:)) / band%dc

    ! Prohibit wavenumber 0.
    desc%src_tau(0) = 0

  end function gaussian_cm_desc

  subroutine gw_cm_src(ncol, band, desc, u, v, frontgf, &
                       src_level, tend_level, tau, ubm, ubi, xv, yv, c)

    use gw_utils, only: get_unit_vector, dot_2d, midpoint_interp

    !-----------------------------------------------------------------------
    ! Driver for multiple gravity wave drag parameterization.
    !
    ! The parameterization is assumed to operate only where water vapor
    ! concentrations are negligible in determining the density.
    !-----------------------------------------------------------------------

    integer, intent(in) :: ncol

    ! Wavelengths triggered by frontogenesis.
    type(GWBand), intent(in) :: band

    ! Bandification of how to produce the gravity wave bandtrum.
    type(CMSourceDesc), intent(in) :: desc

    ! Midpoint zonal/meridional winds.
    real(r8), intent(in) :: u(ncol,pver), v(ncol,pver)
    ! Frontogenesis function.
    real(r8), intent(in) :: frontgf(:,:)

    ! Indices of top gravity wave source level and lowest level where wind
    ! tendencies are allowed.
    integer, intent(out) :: src_level(ncol)
    integer, intent(out) :: tend_level(ncol)

    ! Wave Reynolds stress.
    real(r8), intent(out) :: tau(ncol,-band%ngwv:band%ngwv,pver+1)
    ! Projection of wind at midpoints and interfaces.
    real(r8), intent(out) :: ubm(ncol,pver), ubi(ncol,pver+1)
    ! Unit vectors of source wind (zonal and meridional components).
    real(r8), intent(out) :: xv(ncol), yv(ncol)
    ! Phase speeds.
    real(r8), intent(out) :: c(ncol,-band%ngwv:band%ngwv)

    integer k, l

    ! Whether or not to launch waves in this column.
    logical launch_wave(ncol)

    ! Zonal/meridional wind averaged over source region.
    real(r8) usrc(ncol), vsrc(ncol)

    !------------------------------------------------------------------------
    ! Determine the source layer wind and unit vectors, then project winds.
    !------------------------------------------------------------------------

    ! Just use the source level interface values for the source wind speed
    ! and direction (unit vector).
    src_level = desc%ksrc
    tend_level = desc%ksrc
    usrc = 0.5_r8*(u(:,desc%ksrc+1)+u(:,desc%ksrc))
    vsrc = 0.5_r8*(v(:,desc%ksrc+1)+v(:,desc%ksrc))

    ! Get the unit vector components and magnitude at the surface.
    call get_unit_vector(usrc, vsrc, xv, yv, ubi(:,desc%ksrc+1))

    ! Project the local wind at midpoints onto the source wind.
    do k = 1, desc%ksrc
      ubm(:,k) = dot_2d(u(:,k), v(:,k), xv, yv)
    end do

    ! Compute the interface wind projection by averaging the midpoint winds.
    ! Use the top level wind at the top interface.
    ubi(:,1) = ubm(:,1)

    ubi(:,2:desc%ksrc) = midpoint_interp(ubm(:,1:desc%ksrc))

    !-----------------------------------------------------------------------
    ! Gravity wave sources
    !-----------------------------------------------------------------------

    tau = 0

    ! GW generation depends on frontogenesis at specified level (may be below
    ! actual source level).
    launch_wave = (frontgf(:,desc%kfront) > desc%frontgfc)

    do l = -band%ngwv, band%ngwv
      where (launch_wave)
        tau(:,l,desc%ksrc+1) = desc%src_tau(l)
      end where
    end do

    ! Set phase speeds as reference speeds plus the wind speed at the source
    ! level.
    c = spread(band%cref, 1, ncol) + &
        spread(ubi(:,desc%ksrc+1), 2, 2*band%ngwv+1)

  end subroutine gw_cm_src

end module gw_front
